module Zeno.Core where

import Prelude ()
import StdImports
import Data.IntMap ( IntMap )
import qualified Data.IntMap as IMap
import qualified Data.Map as Map
import qualified Data.Set as Set

import Zeno.Clause
import Zeno.Type
import Zeno.Term
import Zeno.Function
import Zeno.ProofLogic

debug = False

type ZId = Int
type ZTerm = Term ZVar
type ZType = Type ZTypeVar
type ZClause = Clause ZVar
type ZEquality = Equality ZVar
type ZTermSubstitution = TermSubstitution ZVar
type ZTypeSubstitution = TypeSubstitution ZTypeVar
type ZProofStep = ProofStep ZVar
type ZProof = Proof ZVar
type ZFunction = Function ZVar
type ZExpression = Expression ZVar
type ZWhen = When ZVar

data ZVar
  = ZVar        { zVarId :: !ZId,
                  zVarType :: !ZType,
                  zVarClass :: !ZVarClass }
                  
instance Eq ZVar where
  var1 == var2 = zVarId var1 == zVarId var2
  
instance Ord ZVar where
  var1 `compare` var2 = 
    zVarId var1 `compare` zVarId var2

data ZVarClass 
  = ZVariable 
  | ZFunction !ZFunction
  | ZConstructor
  deriving ( Eq )

data ZTypeVar
  = ZDataType   { zTypeVarId :: !ZId,
                  zDataTypeArgs :: ![ZTypeVar],
                  zDataTypeCons :: ![ZVar] }
                  
  | ZPolyType   { zTypeVarId :: !ZId }

instance Show ZTypeVar where
  show (ZDataType id _ _) = showId id
  show (ZPolyType id) = showId id
  
instance Eq ZTypeVar where
  (ZDataType id1 _ _) == (ZDataType id2 _ _) = id1 == id2
  (ZPolyType id1) == (ZPolyType id2) = id1 == id2
  _ == _ = False
  
instance Ord ZTypeVar where
  tv1 `compare` tv2 = zTypeVarId tv1 `compare` zTypeVarId tv2
  
data Zeno
  = Zeno        { zFunctions :: !(Map String ZFunction),
                  zLemmas :: !(Map String ZClause),
                  zDataTypes :: ![ZTypeVar],
                  zNames :: Map ZId String,
                  zCounter :: !ZId }

instance IdCounter Zeno where
  newId zeno = 
    let id = zCounter zeno + 1
    in (id, zeno { zCounter = id })
                  
emptyZeno 
  = Zeno        { zFunctions = mempty,
                  zLemmas = mempty,
                  zDataTypes = mempty,
                  zNames = mempty,
                  zCounter = 1 }
                  
type ZenoState = State Zeno
type ZenoReader = Reader Zeno

instance NFData Zeno where
  rnf (Zeno funs lems dts names ids) =
    rnf funs `seq` rnf lems `seq` rnf dts `seq` 
      rnf names `seq` rnf ids
  
instance NFData ZVar where
  rnf (ZVar id t cls) = rnf id `seq` rnf t
  
instance NFData ZTypeVar where
  rnf (ZDataType id args cons) = rnf id `seq` rnf args
  rnf (ZPolyType id) = rnf id
                                    
polyTypeVar :: ZType -> Bool
polyTypeVar (TypeVar (ZPolyType _)) = True
polyTypeVar _  = False

dataTypeVar :: ZType -> Bool
dataTypeVar (TypeVar (ZDataType {})) = True
dataTypeVar _ = False

functionVar :: ZVar -> Bool
functionVar var =
  case zVarClass var of
    ZFunction _ -> True
    _ -> False

functionTerm :: ZTerm -> Bool
functionTerm = functionVar . termFunction

irreducibleTerm :: ZTerm -> Bool
irreducibleTerm = all functionVar

groundTerm :: ZTerm -> Bool
groundTerm = all ((== ZConstructor) . zVarClass)

unequalTerms :: ZTerm -> ZTerm -> Bool
unequalTerms t1 t2 = 
  constructorTerm t1 &&
  constructorTerm t2 &&
  termFunction t1 /= termFunction t2
  
recursiveType :: ZType -> Bool
recursiveType ztype = 
  let (type_func : _) = flattenAppType ztype
  in  case type_func of
        TypeVar tvar -> recursiveTypeVar tvar
        _ -> False
 
recursiveTypeVar :: ZTypeVar -> Bool
recursiveTypeVar (ZPolyType {}) = False
recursiveTypeVar dt@(ZDataType _ _ cons) =  any (elem dt) $ 
  map (concatMap toList . butlast . flattenFunType . zVarType) cons

variables :: Foldable f => f ZVar -> [ZVar]
variables = nub . filter ((== ZVariable) . zVarClass) . toList

variableTerm :: ZTerm -> Bool
variableTerm term = isVar term &&
  zVarClass (termFunction term) == ZVariable

constructorTerm :: ZTerm -> Bool
constructorTerm term = 
  zVarClass (termFunction term) == ZConstructor
  
getLemma :: MonadReader Zeno m => String -> m ZClause
getLemma name = do
  lemmas <- zLemmas `liftM` ask
  case Map.lookup name lemmas of
    Nothing -> error $ "Lemma \"" ++ name ++ "\" not found"
    Just lemma -> return lemma

freshVariable :: ZVar -> ZenoState ZVar
freshVariable (ZVar id typ cls) = do
  new_id <- newIdS
  return $ ZVar new_id typ cls
  
lookupFunctionName :: String -> ZenoReader ZFunction
lookupFunctionName name = do
  funcs <- zFunctions <$> ask
  case Map.lookup name funcs of
    Just fun -> return fun
    Nothing -> error $ "Couldn't find function " ++ name

-- This function may give an infinite loop for 
--   non-primitive recursive function definitions
applyFunction :: [ZEquality] -> ZTerm -> ZenoReader (Maybe ZExpression)
applyFunction defs term = do
  let (Var func_var) : args = flattenTerm term
  case (zVarClass func_var, null args) of
    (ZFunction func, False) -> do
      let params = map Var $ funcParameters func
      if length params /= length args
        then return Nothing
        else do
          let arg_mapping = Set.fromList $ params `zip` args
              expr = funcExpression func
              mapped_expr = substituteWithin arg_mapping expr
          Just <$> simplifyExpression mapped_expr
          
    _ -> return Nothing
  where
    simplifyExpression :: ZExpression -> ZenoReader ZExpression
    simplifyExpression expr@(Value term) = return expr              
    simplifyExpression expr@(CaseOf term whens) = do  
      let ((Var func) : args) = flattenTerm term
          
          matchWhen (When term' expr') = do
            let ((Var func') : args') = flattenTerm term'
            guard (func' == func) 
            let arg_mapping = Set.fromList $ args' `zip` args
            return $ substituteWithin arg_mapping expr'
      
      -- Multiple matches here indicate a contradiction in the goal antecedents
      simpl <- case mapMaybe matchWhen whens of
        [expr'] -> simplifyExpression expr'
        [] -> do
          maybe_expr <- applyFunction defs term
          case maybe_expr of
            Just (Value term') -> do
              let expr' = CaseOf term' whens
              if constructorTerm term 
                then simplifyExpression expr'
                else return expr'
            _ -> return expr
            
      let expr' = substituteWithin (Set.fromList defs) simpl
      if expr' /= expr
        then simplifyExpression expr'
        else return expr
  
validClause :: ZClause -> Bool
validClause (Clause vars head body) = all validBodyEq body
  where validBodyEq :: ZEquality -> Bool
        validBodyEq (left_term, right_term) = 
          functionTerm left_term && not (functionTerm right_term)
          
termType :: ZTerm -> ZType
termType term =
  let ((Var fun) : args) = flattenTerm term
      fun_types = flattenFunType (zVarType fun)
      res_types = drop (length args) fun_types
  in unflattenFunType res_types
 
functionTypeTerm :: ZTerm -> Bool  
functionTypeTerm = isFunctionType . termType

varTypeTerm :: ZTerm -> Bool
varTypeTerm = isTypeVar . termType

destructibleType :: ZType -> Bool
destructibleType = dataTypeVar . head . flattenAppType

destructibleTerm :: ZTerm -> Bool
destructibleTerm = destructibleType . termType

destructibleVar :: ZVar -> Bool
destructibleVar = destructibleType . zVarType

christen :: MonadState Zeno m => ZId -> String -> m ()
christen id name = do
  zeno <- get
  let names = Map.insert id name (zNames zeno)
  put $ zeno { zNames = names }
  
addDataType :: MonadState Zeno m => ZTypeVar -> m ()
addDataType dtype = dtype `seq`
  modify $ \z -> z { zDataTypes = dtype : (zDataTypes z) }
  {-
addFunctions :: MonadState Zeno m => [ZVar] -> m ()
addFunctions funcs =
  modify $ \z -> z { zFunctions = funcs ++ (zFunctions z) }
  
addVariables :: MonadState Zeno m => [ZVar] -> m ()
addVariables vars =
  modify $ \z -> z { zVariables = vars ++ (zVariables z) }
  -}
addFunction :: ZFunction -> ZenoState ()
addFunction def = do
  let fun_id = zVarId (funcName def)
  fun_name <- readonly (showZId fun_id)
  modify $ \z -> z 
    { zFunctions = Map.insert fun_name def (zFunctions z) }
  
addLemma :: MonadState Zeno m => String -> ZClause -> m ()
addLemma name lemma =
  modify $ \z -> z { zLemmas = Map.insert name lemma (zLemmas z) }
  
showZId :: MonadReader Zeno m => ZId -> m String
showZId id = (fromMaybe (showId id)) `liftM` (lookupName id)

lookupName :: MonadReader Zeno m => ZId -> m (Maybe String)
lookupName id = Map.lookup id `liftM` zNames `liftM` ask 
{-
findFunction :: MonadReader Zeno m => String -> m ZFunction
findFunction name = do
  funs <- zFunctions `liftM` ask
  funs' <- mapKeysM showZVar funs 
  case Map.lookup name funs' of
    Nothing -> error $ "Function \"" ++ name ++ "\" not found"
    Just fn -> return fn
  -}  
findDataType :: MonadReader Zeno m => String -> m ZTypeVar
findDataType name = do
  dtypes <- zDataTypes `liftM` ask
  names <- mapM (showZId . zTypeVarId) dtypes
  case lookup name (names `zip` dtypes) of
    Nothing -> error $ "Datatype \"" ++ name ++ "\" not found"
    Just dt -> return dt
    
instantiate :: ZTypeSubstitution -> ZVar -> ZenoState (ZTerm, [ZVar])
instantiate bindings var = do
  let var_type = substituteWithin bindings (zVarType var)
      var_types = flattenFunType var_type
      (arg_types, [res_type]) = splitAt (length var_types - 1) var_types
  new_ids <- replicateM (length arg_types) newIdS
  let args' = zipWith ($) (map ZVar new_ids) arg_types
      args = map ($ ZVariable) args'
      vars = filter ((== res_type) . zVarType) args
      term = unflattenTerm (map Var (var : args))
  return (term, vars)
  
applyTermDefinitions :: WithinTraversable ZTerm (t ZVar) => 
    [ZEquality] -> t ZVar -> ZenoReader ((t ZVar), Bool)
applyTermDefinitions defs t = do
  let mapped_t = mapWithinM applyTermDefinition t
  (t', any) <- runWriterT mapped_t
  if getAny any 
    then do
      (t'', _) <- applyTermDefinitions defs t'
      return (t'', True)
    else 
      return (t', False)
  where
    applyTermDefinition :: ZTerm -> WriterT Any ZenoReader ZTerm
    applyTermDefinition term = do
      expr <- lift $ applyFunction defs term
      case expr of 
        Just (Value term') -> tell (Any True) >> return term'
        _ -> return term
    
getName :: MonadReader Zeno m => ZId -> m String
getName = (liftM fromJust) . lookupName

largestId :: ZClause -> ZId
largestId = maximum . (map zVarId) . toList
 
clausifyFunction :: ZFunction -> [ZClause]
clausifyFunction (Function _ _ (Value (Var (ZVar _ _ (ZFunction func))))) =
  clausifyFunction func
clausifyFunction (Function _ _ expr) =
  map normaliseClause (clausifyExpression expr)
  where
    clausifyExpression :: ZExpression -> [ZClause]
    clausifyExpression (Value (App (App (Var _) x) y)) = 
      [ Clause [] (x, y) [] ]
    clausifyExpression (Value _) = []
    clausifyExpression (CaseOf ct ws) = concatMap clausifyWhen ws
      where
        clausifyWhen (When wt expr) = 
          map (addCondition (ct, wt)) (clausifyExpression expr)
          
    normaliseClause :: ZClause -> ZClause
    normaliseClause cls@(Clause vars eq conds) =
      case find (variableTerm . fst) conds of
        Just (from, to) ->
          let replace = replaceWithin from to
              conds' = map (mapPair replace) conds
              conds'' = filter (not . uncurry equalTerms) conds'
          in normaliseClause (Clause vars eq conds'')
        Nothing -> 
          let (conds', any) = runWriter $ concatMapM factoriseEquality conds
          in  if getAny any 
                then normaliseClause (Clause vars eq conds')
                else cls
      where
        factoriseEquality :: ZEquality -> Writer Any [ZEquality]
        factoriseEquality eq@(t1, t2) = 
          let ((Var func1) : args1) = flattenTerm t1
              ((Var func2) : args2) = flattenTerm t2
          in  if func1 == func2 && zVarClass func1 == ZConstructor
                then tell (Any True) >> return (args1 `zip` args2)
                else return [eq]
            
showZTerm :: MonadReader Zeno m => ZTerm -> m String
showZTerm = liftM showUnquoted . mapM showZVar

showZVar :: MonadReader Zeno m => ZVar -> m String
showZVar = showZId . zVarId

showZVarTyped :: MonadReader Zeno m => ZVar -> m String
showZVarTyped var = do
  var_s <- showZVar var
  type_s <- showZType (zVarType var)
  return (var_s ++ " :: " ++ type_s)

showZType :: MonadReader Zeno m => ZType -> m String
showZType = liftM showUnquoted . mapM showZTypeVar

showZTypeVar :: MonadReader Zeno m => ZTypeVar -> m String
showZTypeVar = showZId . zTypeVarId
 
showFullZTypeVar :: MonadReader Zeno m => ZTypeVar -> m String
showFullZTypeVar (ZPolyType id) = showZId id
showFullZTypeVar dtype@(ZDataType id args cons) = do
  name_s <- showZId id
  args_s <- intercalate " " `liftM` mapM showZTypeVar args
  cons_s <- intercalate "\n  " `liftM` mapM showZVarTyped cons
  let rec_s = if recursiveTypeVar dtype then "rec " else ""
  return $ rec_s ++ "type " ++ name_s ++ " " 
    ++ args_s ++ "\n  " ++ cons_s
     
showZEquality :: MonadReader Zeno m => ZEquality -> m String
showZEquality (t1, t2) = do
  t1' <- showZTerm t1
  t2' <- showZTerm t2
  return $ showEquality (t1', t2')
    
showZClause :: MonadReader Zeno m => ZClause -> m String
showZClause = (liftM showUnquoted) . mapM showZVar

showZTermSub :: MonadReader Zeno m => ZTermSubstitution -> m String
showZTermSub = (liftM show) . mapM (mapPairM showZTerm) . toList

showZFunction :: MonadReader Zeno m => ZFunction -> m String
showZFunction = (liftM showUnquoted) . mapM showZVar

showZExpression :: MonadReader Zeno m => ZExpression -> m String
showZExpression = (liftM showUnquoted) . mapM showZVar

showZProof :: ZProof -> ZenoReader String
showZProof proof = do
  let ([proof'], inds) = runWriter $ findInducts True proof
  proof_s <- showUnquoted <$> mapM showZVar proof'
  inds_s <- intercalate "\n        " <$> 
    mapM showZClause (map freeAllVariables inds)
  return $ proof_s ++ "\n\nProven: " ++ inds_s
  where
    freeAllVariables :: ZClause -> ZClause
    freeAllVariables cls = cls { clauseVariables = variables cls }
  
    findInducts :: Bool -> ZProof -> Writer [ZClause] [ZProof]
    findInducts _ Unknown = return [Unknown]
    findInducts _ ds@(Disproven {}) = return [ds] 
    findInducts ind pr@(Proven step@(ProofStep _ Nothing []) ps) = do
      ps' <- concatMapM (findInducts True) ps
      return [Proven step ps']
    findInducts ind (Proven (ProofStep "fct" _ []) ps) =
      concatMapM (findInducts ind) ps
    findInducts ind (Proven step@(ProofStep name (Just cls) _) ps) = do
      let ind_step = name == "ind" || name == "cse"
          ind' = not ind_step && 
            ("hyp" `isPrefixOf` name || 
             "gen" `isPrefixOf` name || ind)
      ps' <- concatMapM (findInducts ind') ps
      when (ind_step && ind) (tell [cls])
      if ind_step
        then return ps'
        else return [Proven step ps']
      

instance Show ZVar where
  show = show . zVarId

instance Show Zeno where
  show zeno = runReader reader zeno
    where
      datatypes = zDataTypes zeno
      functions = Map.elems $ zFunctions zeno
      names = Map.toList $ zNames zeno
      lemmas = Map.toList $ zLemmas zeno
      
      reader :: MonadReader Zeno m => m String
      reader = do
        datatypes' <- mapM showFullZTypeVar datatypes
        let datatypes_s = intercalate "\n\n" datatypes'
        
        functions' <- mapM (mapM showZVar) functions
        let functions_s = intercalate "\n\n" (map showUnquoted functions')
        
        let showLemma (name, lemma) = do
              lemma' <- showZClause lemma
              return $ name ++ "\n  " ++ lemma'
        lemmas' <- mapM showLemma lemmas
        let lemmas_s = intercalate "\n\n" lemmas'
        
        return $ "\n\n" ++ datatypes_s ++ 
          "\n\n" ++ functions_s ++ 
          "\n\n" ++ lemmas_s

          
-- Debug section
          
debugZeno :: IORef Zeno
debugZeno = unsafePerformIO (newIORef emptyZeno)

traceZeno :: String -> ZenoReader String -> b -> b
traceZeno prefix reader expr
  | debug = unsafePerformIO $ do
      zeno <- readIORef debugZeno
      putStrLn $ "\n" ++ prefix
      putStrLn $ runReader reader zeno
      putStrLn ""
      return expr
  | otherwise = expr

traceZClause :: String -> ZClause -> a -> a
traceZClause prefix clause = 
  traceZeno prefix $ showZClause clause
  
traceZVar prefix var = 
  traceZeno prefix $ showZVar var
  
traceZClauses :: String -> [ZClause] -> a -> a
traceZClauses str = flip $ foldr (traceZClause str)
  
traceZEqualities :: String -> [ZEquality] -> a -> a
traceZEqualities str = flip $ foldr (traceZEquality str)
  
traceZEquality :: String -> ZEquality -> a -> a
traceZEquality s eq = traceZClause s (Clause mempty eq mempty)

traceZTerm :: String -> ZTerm -> a -> a
traceZTerm str term = traceZeno str (showZTerm term) 

traceZTerms :: String -> [ZTerm] -> a -> a
traceZTerms str = flip $ foldr (traceZTerm str)

traceZType :: ZType -> a -> a
traceZType t = traceZeno "" (showZType t) 

traceZTypes :: [ZType] -> a -> a
traceZTypes = flip $ foldr traceZType

traceZSub :: String -> ZTermSubstitution -> a -> a
traceZSub pref sub = traceZeno pref $ do
  sub' <- mapM showZEquality (toList sub)
  return $ intercalate "\n" sub'
  
traceZ :: String -> a -> a
traceZ 
  | debug = trace
  | otherwise = flip const

